package cn.iocoder.yudao.framework.mybatis.autotable.creator;

import cn.hutool.core.io.IoUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.iocoder.yudao.framework.mybatis.autotable.ddl.DDLGenerator;
import cn.iocoder.yudao.framework.mybatis.autotable.scanner.TableScanner;
import com.baomidou.mybatisplus.annotation.TableName;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;

import javax.sql.DataSource;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.nio.charset.StandardCharsets;
import java.sql.*;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

@Slf4j
public class TableAutoCreator implements ApplicationRunner {

    private final DataSource dataSource;
    private final DDLGenerator ddlGenerator;
    private final String[] basePackages;

    public TableAutoCreator(DataSource dataSource, DDLGenerator ddlGenerator, String[] basePackages) {
        this.dataSource = dataSource;
        this.ddlGenerator = ddlGenerator;
        this.basePackages = basePackages;
    }

    @Override
    public void run(ApplicationArguments args) {
        Set<Class<?>> entities = TableScanner.scanEntities(basePackages);
        log.info("Scanned {} entities with @TableName", entities.size());
        Set<String> success = new HashSet<>();
        Set<String> failed = new HashSet<>();
        AtomicInteger index = new AtomicInteger(0);
        File file = resolveOutputDir();
        String tag = SpringUtil.getProperty("spring.autotable.tag");
        boolean tenantEnable = Boolean.parseBoolean(SpringUtil.getProperty("spring.autotable.tenant-enable"));
        log.info("write ddl to file {}", file.getName());
        try (Connection conn = dataSource.getConnection()) {
            Map<String, Class<?>> tableNames = entities.stream().collect(Collectors.toMap(this::resolveTableName, entity -> entity));
            Map<String, Boolean> existRes = checkTablesExist(conn, tableNames.keySet());
            long count = existRes.values().stream().filter(e->!e).count();
            existRes.forEach((tableName, exists) -> {
                if (!exists) {
                    log.warn("扫描结果:Table {} does not exist, creating...", tableName);
                    try {
                        log.info("开始执行表 {} 的创建语句...第 {}/{}", tableName,  index.addAndGet(1), count);
                        boolean result = createTable(conn, tableNames.get(tableName), tableName, tag, tenantEnable, file);
                        if (result){
                            success.add(tableName);
                        }else{
                            failed.add(tableName);
                        }
                    } catch (SQLException e) {
                        throw new RuntimeException(e);
                    }
                }
            });

            log.warn("在{}下共计扫描实体类{}个,共计缺失 {} 个表,已完成自动创建,成功:{},失败:{}个(明细:{})", basePackages, entities.size(), count
                    , success.size(), failed.size(), failed);

        } catch (SQLException e) {
            throw new RuntimeException("Auto table creation failed", e);
        }
    }

    private File resolveOutputDir() {
        String outDir = SpringUtil.getProperty("spring.autotable.outdir");
        String prefix = System.getProperty("user.dir");
        outDir = outDir.replace("/", File.separator);
        String dir = prefix + File.separator + outDir;
        File file = new File(dir);
        log.info("output dir: {}", dir);
        if (!file.exists()) {
            try {
                boolean mkDirs = file.createNewFile();
                if (!mkDirs) {
                    log.error("create dir {} error", dir);
                    throw new RuntimeException("create dir " + dir + " error");
                }
                log.info("create dir {} success", dir);
                return file;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        //如果存在的话,则先清除文件内容
        else {
            try {
                FileOutputStream fileOutputStream = new FileOutputStream(file);
                IoUtil.write(fileOutputStream, true, new byte[0]);
                return file;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

    private String resolveTableName(Class<?> entity) {
        TableName annotation = entity.getAnnotation(TableName.class);
        return annotation != null ? annotation.value() :
                StringUtils.camelToUnderline(entity.getSimpleName());
    }


    private Map<String, Boolean> checkTablesExist(Connection conn, Set<String> tableNames) throws SQLException {
        DatabaseMetaData meta = conn.getMetaData();
        Map<String, Boolean> result = new HashMap<>();
        try (ResultSet rs = meta.getTables(null, null, "%", new String[]{"TABLE"})) {
            // 获取所有表名集合
            Set<String> existingTables = new HashSet<>();
            while (rs.next()) {
                existingTables.add(rs.getString("TABLE_NAME").toLowerCase());
            }
            // 批量比对
            for (String tableName : tableNames) {
                result.put(tableName, existingTables.contains(tableName.toLowerCase()));
            }
        }
        return result;
    }


    private boolean createTable(Connection conn, Class<?> entity, String tableName, String tag, boolean tenantEnable, File outDir) throws SQLException {
        String ddl = ddlGenerator.generateDDL(dataSource, entity, tag, tenantEnable);
        if (outDir.exists()) {
            //写入 outDir
            try {
                FileOutputStream fileOutputStream = new FileOutputStream(outDir, true);
                IoUtil.write(fileOutputStream, false, ddl.getBytes(StandardCharsets.UTF_8));
            } catch (FileNotFoundException e) {
                log.error("write ddl to {} error:", outDir);
            }
        }
        try (Statement stmt = conn.createStatement()) {
            String[] split = ddl.split(";");
            for (String s : split) {
                stmt.addBatch(s);
            }
            stmt.executeBatch();
            return true;
        } catch (Exception e) {
            log.error("auto create statement {}  error :(", tableName, e);
            return false;
        }
    }
}